"""
RAG文档摄取模块 - 文档上传、分块、嵌入和存储

这个模块提供了一个完整的文档处理流水线，支持多种文档格式，
自动完成文档加载、文本分块、向量嵌入和向量数据库存储。

支持的文档格式：PDF, PPTX, TXT, DOCX
嵌入模型：BAAI/bge-m3 多语言嵌入模型
向量数据库：Chroma（支持本地持久化）

主要特性：
  - 增量添加模式：向现有数据库添加新文档，不会丢失原有数据（默认模式）
  - 覆盖重建模式：删除现有数据库并重新创建（需谨慎使用）

使用示例：
  # 增量添加新文档（推荐）
  vector_db = ingest_documents(
      file_paths=["new_doc1.txt", "new_doc2.pdf"],
      mode="append"  # 默认值
  )
  
  # 覆盖重建整个数据库
  vector_db = ingest_documents(
      file_paths=["doc1.txt", "doc2.pdf"],
      mode="overwrite"  # 谨慎使用
  )

作者：AI Assistant
日期：2025-09-19
更新：2025-10-12 - 添加增量添加功能
"""

import os
import logging
from pathlib import Path
from typing import List, Optional

# Langchain核心导入
from langchain.schema import Document
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.embeddings.base import Embeddings

# Langchain文档加载器
from langchain_community.document_loaders import (
    PyPDFLoader,
    # UnstructuredPowerPointLoader,  # 暂时注释掉PPT支持
    TextLoader,
    #UnstructuredWordDocumentLoader
)

# HuggingFace嵌入模型（已弃用，改用API调用）
# from langchain_community.embeddings import HuggingFaceEmbeddings
# BGE-M3 API Embeddings
from bge_m3_api_embeddings import BGEM3APIEmbeddings

# Chroma向量数据库
from langchain_community.vectorstores import Chroma

# 设置日志
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


def load_documents(file_paths: List[str]) -> List[Document]:
    """
    根据文件类型加载文档，返回Document对象列表。
    
    Args:
        file_paths: 文件路径列表
        
    Returns:
        List[Document]: 加载的文档对象列表
        
    Raises:
        FileNotFoundError: 当文件不存在时
        ValueError: 当文件类型不支持时
    """
    documents = []
    
    # 定义支持的文件类型和对应的加载器
    loaders_map = {
        '.pdf': PyPDFLoader,
        # '.pptx': UnstructuredPowerPointLoader,  # 暂时注释掉PPT支持
        '.txt': TextLoader,
        #'.docx': UnstructuredWordDocumentLoader
    }
    
    for file_path in file_paths:
        # 检查文件是否存在
        if not os.path.exists(file_path):
            logger.error(f"文件不存在: {file_path}")
            raise FileNotFoundError(f"文件不存在: {file_path}")
        
        # 获取文件扩展名
        file_extension = Path(file_path).suffix.lower()
        
        # 检查文件类型是否支持
        if file_extension not in loaders_map:
            supported_types = ', '.join(loaders_map.keys())
            error_msg = f"不支持的文件类型: {file_extension}。支持的类型: {supported_types}"
            logger.error(error_msg)
            raise ValueError(error_msg)
        
        try:
            # 获取对应的加载器类
            loader_class = loaders_map[file_extension]
            
            # 创建加载器实例并加载文档
            if file_extension == '.txt':
                # TextLoader需要指定编码
                loader = loader_class(file_path, encoding='utf-8')
            else:
                loader = loader_class(file_path)
            
            # 加载文档
            file_documents = loader.load()
            
            # 为每个文档添加元数据
            for doc in file_documents:
                doc.metadata.update({
                    'source_file': file_path,
                    'file_type': file_extension
                })
            
            documents.extend(file_documents)
            logger.info(f"成功加载文档: {file_path} (共{len(file_documents)}个文档)")
            
        except Exception as e:
            logger.error(f"加载文档失败: {file_path}, 错误: {str(e)}")
            raise
    
    logger.info(f"总共加载了 {len(documents)} 个文档")
    return documents


def split_documents(documents: List[Document]) -> List[Document]:
    """
    将文档分块处理，便于向量化和检索。
    
    Args:
        documents: 原始文档列表
        
    Returns:
        List[Document]: 分块后的文档列表
    """
    # 创建文本分割器
    text_splitter = RecursiveCharacterTextSplitter(
        chunk_size=1000,        # 每个分块最大字符数
        chunk_overlap=200,      # 分块之间的重叠字符数
        length_function=len,    # 计算长度的函数
        separators=['\n\n', '\n', ' ', '']  # 分割优先级
    )
    
    try:
        # 分割文档
        chunks = text_splitter.split_documents(documents)
        
        # 为分块添加额外的元数据
        for i, chunk in enumerate(chunks):
            chunk.metadata.update({
                'chunk_id': i,
                'chunk_size': len(chunk.page_content)
            })
        
        logger.info(f"文档分块完成，共生成 {len(chunks)} 个分块")
        return chunks
        
    except Exception as e:
        logger.error(f"文档分块失败: {str(e)}")
        raise


def get_optimal_device():
    """
    获取最优设备配置，自动检测GPU可用性
    
    Returns:
        str: 推荐的设备类型 ('cuda', 'mps', 'cpu')
    """
    try:
        import torch
        
        # 检查CUDA
        if torch.cuda.is_available():
            # 获取GPU内存信息
            try:
                memory_total = torch.cuda.get_device_properties(0).total_memory / 1024**2
                memory_allocated = torch.cuda.memory_allocated() / 1024**2
                memory_free = memory_total - memory_allocated
                
                # 至少需要2GB GPU内存来运行BGE-M3
                if memory_free > 2048:  
                    device_name = torch.cuda.get_device_name()
                    logger.info(f"🚀 检测到可用GPU: {device_name}")
                    logger.info(f"GPU内存: {memory_free:.0f}MB 可用 / {memory_total:.0f}MB 总计")
                    return 'cuda'
                else:
                    logger.warning(f"⚠️  GPU内存不足({memory_free:.0f}MB)，降级到CPU")
            except Exception as e:
                logger.warning(f"⚠️  GPU内存检查失败: {e}，降级到CPU")
        
        # 检查Apple Silicon MPS
        if hasattr(torch.backends, 'mps') and torch.backends.mps.is_available():
            logger.info("🍎 检测到Apple Silicon GPU，使用MPS加速")
            return 'mps'
        
        logger.info("💻 使用CPU进行计算")
        return 'cpu'
        
    except ImportError:
        logger.warning("⚠️  PyTorch未安装，使用CPU")
        return 'cpu'
    except Exception as e:
        logger.warning(f"⚠️  设备检测失败: {e}，使用CPU")
        return 'cpu'


def create_embeddings_model() -> Embeddings:
    """
    创建BAAI/bge-m3嵌入模型实例，通过API调用。
    
    BGE-M3是北京人工智能研究院开发的多语言嵌入模型，支持中文、英文等多种语言，
    通过ModelArts MaaS API调用，无需本地加载模型，向量维度为1024。
    
    Returns:
        Embeddings: 嵌入模型实例
        
    Raises:
        Exception: 当模型初始化失败时
    """
    try:
        logger.info("初始化BGE-M3 API嵌入模型...")
        
        # 创建BGE-M3 API嵌入模型实例
        embeddings_model = BGEM3APIEmbeddings()
        
        logger.info("✅ 成功初始化BAAI/bge-m3嵌入模型（API调用）")
        logger.info("模型特性: 多语言支持、1024维向量、API调用")
        logger.info("⚡ API调用模式：无需本地模型，无需GPU")
        
        return embeddings_model
        
    except Exception as e:
        logger.error(f"❌ 创建嵌入模型失败: {str(e)}")
        logger.error("请检查:")
        logger.error("1. 是否已设置环境变量 BGE_M3_API_KEY 或 MAAS_API_KEY")
        logger.error("2. 是否已设置环境变量 BGE_M3_API_URL（可选，有默认值）")
        logger.error("3. 网络连接是否正常")
        raise


def store_documents_in_vector_db(
    chunks: List[Document], 
    embeddings_model: Embeddings, 
    persist_directory: str,
    mode: str = "append"
) -> Chroma:
    """
    将文档分块和嵌入向量存储到Chroma向量数据库。
    
    Args:
        chunks: 文档分块列表
        embeddings_model: 嵌入模型实例
        persist_directory: 数据库持久化目录
        mode: 存储模式，"append"（增量添加）或"overwrite"（覆盖重建），默认为"append"
        
    Returns:
        Chroma: Chroma向量数据库实例
    """
    try:
        # 确保持久化目录存在
        os.makedirs(persist_directory, exist_ok=True)
        
        # 检查数据库是否已存在
        db_exists = os.path.exists(os.path.join(persist_directory, "chroma.sqlite3"))
        
        if mode == "overwrite" or not db_exists:
            # 覆盖模式或数据库不存在：创建新的向量数据库
            if db_exists and mode == "overwrite":
                logger.warning("⚠️  覆盖模式：将删除现有数据库并重新创建")
                # 可选：删除旧数据库文件（Chroma会自动处理）
            
            vector_db = Chroma.from_documents(
                documents=chunks,
                embedding=embeddings_model,
                persist_directory=persist_directory,
                collection_name="rag_documents"
            )
            
            logger.info(f"✅ {'重新创建' if db_exists else '创建'}向量数据库，存储了 {len(chunks)} 个文档分块")
            
        else:
            # 增量添加模式：加载现有数据库并添加新文档
            logger.info("📂 检测到现有数据库，使用增量添加模式")
            
            # 加载现有向量数据库
            vector_db = Chroma(
                persist_directory=persist_directory,
                embedding_function=embeddings_model,
                collection_name="rag_documents"
            )
            
            # 获取添加前的文档数量
            collection = vector_db._collection
            old_count = collection.count()
            logger.info(f"📊 数据库现有 {old_count} 个文档分块")
            
            # 增量添加新文档
            vector_db.add_documents(chunks)
            
            # 获取添加后的文档数量
            new_count = collection.count()
            added_count = new_count - old_count
            
            logger.info(f"✅ 成功添加 {added_count} 个新文档分块")
            logger.info(f"📊 数据库现共有 {new_count} 个文档分块")
        
        logger.info(f"💾 数据库持久化路径: {persist_directory}")
        
        return vector_db
        
    except Exception as e:
        logger.error(f"❌ 存储到向量数据库失败: {str(e)}")
        raise


def ingest_documents(
    file_paths: List[str], 
    persist_directory: str = "./chroma_db",
    mode: str = "append"
) -> Chroma:
    """
    文档摄取主函数 - 编排整个文档处理流水线。
    
    这是整个模块的入口点，完成从文档加载到向量存储的全过程。
    
    Args:
        file_paths: 要处理的文件路径列表
        persist_directory: 向量数据库持久化目录，默认为"./chroma_db"
        mode: 存储模式，"append"（增量添加）或"overwrite"（覆盖重建），默认为"append"
              - "append": 向现有数据库添加新文档（推荐，不会丢失原有数据）
              - "overwrite": 删除现有数据库并重新创建（谨慎使用）
        
    Returns:
        Chroma: 填充了文档的向量数据库实例
        
    Raises:
        Exception: 任何步骤失败时抛出异常
    """
    logger.info("=" * 60)
    logger.info("开始文档摄取流程...")
    logger.info(f"📂 待处理文件: {file_paths}")
    logger.info(f"💾 数据库路径: {persist_directory}")
    logger.info(f"🔧 存储模式: {mode} {'(增量添加)' if mode == 'append' else '(覆盖重建)'}")
    logger.info("=" * 60)
    
    try:
        # 第一步：加载文档
        logger.info("步骤 1/4: 加载文档...")
        documents = load_documents(file_paths)
        
        # 第二步：分块处理
        logger.info("步骤 2/4: 文档分块...")
        chunks = split_documents(documents)
        
        # 第三步：创建嵌入模型
        logger.info("步骤 3/4: 初始化嵌入模型...")
        embeddings_model = create_embeddings_model()
        
        # 第四步：存储到向量数据库
        logger.info(f"步骤 4/4: 存储到向量数据库（{mode}模式）...")
        vector_db = store_documents_in_vector_db(chunks, embeddings_model, persist_directory, mode)
        
        logger.info("=" * 60)
        logger.info("✅ 文档摄取流程完成！")
        logger.info("=" * 60)
        return vector_db
        
    except Exception as e:
        logger.error(f"❌ 文档摄取流程失败: {str(e)}")
        raise


# 示例用法和测试代码
if __name__ == "__main__":
    """
    示例用法：演示如何使用文档摄取模块
    
    运行前请确保：
    1. 安装必要的依赖包：
       pip install langchain langchain-community sentence-transformers chromadb
       pip install unstructured[all-docs] pypdf python-docx
    
    2. 准备示例文件（可选）：
       创建./data目录并放置示例文件
    
    注意：BGE-M3模型使用API调用，需要配置API密钥
    """
    
    # BGE-M3模型通过API调用，无需本地模型
    print("🚀 使用BAAI/bge-m3嵌入模型（API调用模式，无需本地模型）")
    
    # 示例文件路径（请根据实际情况修改）
    example_files = [
        "./data/ai_overview.txt",
        # "./data/presentation.pptx",  # 暂时注释掉PPT支持
        "./data/notes.txt",
        "./data/outline.pdf",
        "./data/tech_guide.txt",
        "./data/大数据技术原理与应用(第三版) (林子雨) 55885.pdf",
        "./data/TML.pdf"
    ]
    
    # 检查示例文件是否存在
    existing_files = [f for f in example_files if os.path.exists(f)]
    
    if not existing_files:
        print("📁 未找到示例文件，请创建以下目录和文件进行测试：")
        print()
        print("或者修改example_files列表，指向您的实际文件。")
        exit(1)
    
    try:
        # 检查是否已存在数据库
        db_exists = os.path.exists("./chroma_db/chroma.sqlite3")
        
        if db_exists:
            print("\n" + "=" * 60)
            print("📂 检测到现有向量数据库")
            print("=" * 60)
            print("\n请选择运行模式：")
            print("  1. 增量添加模式 (append) - 向现有数据库添加新文档【推荐】")
            print("  2. 覆盖重建模式 (overwrite) - 删除现有数据并重新创建【谨慎】")
            print("  3. 取消操作")
            print()
            
            choice = input("请输入选项 (1/2/3): ").strip()
            
            if choice == "1":
                mode = "append"
                print("\n✅ 已选择：增量添加模式")
            elif choice == "2":
                confirm = input("⚠️  确认要删除现有数据吗？输入 'yes' 确认: ").strip().lower()
                if confirm == "yes":
                    mode = "overwrite"
                    print("\n⚠️  已选择：覆盖重建模式")
                else:
                    print("❌ 已取消操作")
                    exit(0)
            else:
                print("❌ 已取消操作")
                exit(0)
        else:
            mode = "overwrite"  # 首次创建使用overwrite模式
            print("\n📂 未检测到现有数据库，将创建新的向量数据库")
        
        print("\n🚀 开始处理文档...")
        print(f"📂 找到的文件: {existing_files}")
        
        # 执行文档摄取
        vector_store = ingest_documents(
            file_paths=existing_files,
            persist_directory="./chroma_db",
            mode=mode
        )
        
        # 显示结果
        collection = vector_store._collection
        total_count = collection.count()
        
        print(f"\n✅ 成功完成文档摄取！")
        print(f"📊 向量数据库统计信息:")
        print(f"   - 持久化路径: ./chroma_db")
        print(f"   - 集合名称: rag_documents")
        print(f"   - 文档分块总数: {total_count}")
        print(f"   - 运行模式: {mode}")
        
        # 简单测试检索功能
        print(f"\n🔍 测试检索功能...")
        results = vector_store.similarity_search("测试查询", k=3)
        print(f"   - 检索到 {len(results)} 个相关文档分块")
        
        if results:
            print(f"   - 第一个结果预览: {results[0].page_content[:100]}...")
        
        print(f"\n🎉 RAG文档摄取模块测试完成！")
        print(f"\n💡 使用提示:")
        print(f"   - 下次运行时可以继续添加新文档（增量模式）")
        print(f"   - 也可以选择重新构建整个数据库（覆盖模式）")
        
    except Exception as e:
        print(f"❌ 处理失败: {str(e)}")
        exit(1)
